#!/usr/bin/env python

# Copyright 2010 Noam Yorav-Raphael
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# expensefs - show gnucash expenses as a filesystem, so that disk usage
# analyzers will help you explore your expenses!

import sys
import os
import re
import gzip
from datetime import datetime
import xml.etree.cElementTree as ET
import argparse
from stat import S_IFDIR, S_IFREG
from errno import ENOENT
import signal

from fuse import FUSE, FuseOSError, Operations, LoggingMixIn
from record import Record

ns_cache = {}

def ns(s):
    try:
        return ns_cache[s]
    except KeyError:
        ns_cache[s] = r = re.sub(
            r'\b(gnc|act|trn|ts|split):', r'{http://www.gnucash.org/XML/\1}', s)
        return r

class Account(Record):
    def __init__(self, name, guid, subaccs, tras):
        Record.__init__(self, locals())

class Transaction(Record):
    def __init__(self, desc, amount, date):
        Record.__init__(self, locals())

def parse(fn):
    """
    Get a filename, return the root expenses account.
    """
    et = ET.parse(gzip.open(fn))
    exp_accounts = [x for x in et.getiterator(ns('gnc:account'))
                    if x.findtext(ns('act:type')) in ('EXPENSE', 'ROOT')]
    accounts = [Account(x.findtext(ns('act:name')), x.findtext(ns('act:id')), [], [])
                for x in exp_accounts]
    guids = dict((acc.guid, acc) for acc in accounts)
    for i, x in enumerate(exp_accounts):
        parent_guid = x.findtext(ns('act:parent'))
        try:
            parent = guids[parent_guid]
        except KeyError:
            # We ignore expense accounts which are not children of expense accounts
            pass
        else:
            parent.subaccs.append(accounts[i])
    root_i = [i for i, x in enumerate(exp_accounts)
              if x.findtext(ns('act:type')) == 'ROOT'][0]
    # We assume that there's only one root expense account
    root_exp = accounts[root_i].subaccs[0]
    
    for tra in et.getiterator(ns('gnc:transaction')):
        splits = tra.find(ns('trn:splits'))
        for split in splits:
            guid = split.findtext(ns('split:account'))
            try:
                account = guids[guid]
            except KeyError:
                continue
            
            desc = tra.findtext(ns('trn:description'))
            date_s = tra.findtext(ns('trn:date-posted/ts:date'))
            date = datetime.strptime(date_s[:-6], '%Y-%m-%d %H:%M:%S')
            quantity = split.findtext(ns('split:quantity'))
            assert quantity.endswith('/100')
            amount = int(round(int(quantity[:-4])/100.))
            account.tras.append(Transaction(desc, amount, date))
        
    return root_exp

def make_hierarchy(root):
    """
    Get the root expense account. Create a dir hierarchy.
    The hierarchy is a dict which maps names to either dicts (dirs) or numbers
    (size of file/transaction).
    Every transaction appears 3 times, because we have a total, filtering by year
    and filtering by month. It goes like this:
    all/
        (accounts)
    2009-all/
        (accounts)
    2009/
        01/
            (accounts)
        02/
            (accounts)
        ...
    ...
    """
    paths = []
    make_hierarchy_rec(paths, (), root)
    d = {}
    for path, name, size in paths:
        curdict = d
        for element in path:
            curdict = curdict.setdefault(element, {})
        curdict[name] = size
    return d

def make_hierarchy_rec(paths, path, acc):
    """
    Recursive helper for make_hierarchy.
    path - a tuple with names up to acc.name
    paths - a list to which paths are added.
            Every element is a tuple with the path, the name, and size.
    """
    for subacc in acc.subaccs:
        make_hierarchy_rec(paths, path+(subacc.name,), subacc)
    for tra in acc.tras:
        name = '%s (%s)' % (tra.desc, tra.date.strftime('%d-%m-%y'))
        paths.append((('all',)+path, name, tra.amount))
        paths.append((('%d-all'%tra.date.year,)+path, name, tra.amount))
        paths.append((('%d'%tra.date.year, '%02d'%tra.date.month)+path,
                      name, tra.amount))

def flatten(hierarchy):
    """
    Take a hierarchy, return a dict mapping full paths to either dicts (which
    map names to dicts or ints) or ints.
    """
    d = {'/': hierarchy}
    flatten_rec(hierarchy, d, '/')
    return d

def flatten_rec(hierarchy, d, prefix):
    """
    Recursive helper for flatten.
    hierarchy - a dict to read from.
    d - a dict to fill.
    prefix - current prefix.
    """
    for name, val in hierarchy.iteritems():
        name = name.encode('utf8')
        d[prefix+name] = val
        if not isinstance(val, int):
            flatten_rec(val, d, prefix+name+'/')

def make_dirs(dir, hierarchy):
    join = os.path.join
    for name, val in hierarchy.iteritems():
        subdir = join(dir, name.replace('/', "|"))
        os.mkdir(subdir)
        if isinstance(val, int):
            f = open(join(subdir, 'bytes'), 'w')
            f.write('a'*val)
            f.close()
        else:
            make_dirs(subdir, val)

# Add LoggingMixIn as a base class to log.
class Filesystem(Operations):
    def __init__(self, fn):
        """
        fn - filename of the gnucash file.
        """
        self.fn = fn
        self.last_mtime = 0
        # Flattened hierarchy. Updated when gnucash file was updated.
        self.flat = None
        # Attributes which are common to every node. Updated by check_cache.
        self.common_attrs = None
    
    def check_cache(self):
        cur_mtime = os.path.getmtime(self.fn)
        if cur_mtime != self.last_mtime:
            print >> sys.stderr, "Reading gnucash file...",
            root = parse(self.fn)
            hierarchy = make_hierarchy(root)
            self.flat = flatten(hierarchy)
            self.common_attrs = dict(
                st_mtime=cur_mtime, st_ctime=cur_mtime, st_atime=cur_mtime,
                st_uid=os.getuid(), st_gid=os.getgid())
                
            self.last_mtime = cur_mtime
            print >> sys.stderr, "Done."
    
    def readdir(self, path, fh):
        self.check_cache()
        
        try:
            val = self.flat[path]
        except KeyError:
            raise FuseOSError(ENOENT)
        
        if isinstance(val, int):
            return ['.', '..', 'bytes']
        else:
            return ['.', '..']+[name.encode('utf8') for name in val.iterkeys()]
    
    def getattr(self, path, fh=None):
        self.check_cache()
        
        if path.endswith('/bytes'):
            try:
                size = self.flat[path[:-6]]
            except KeyError:
                raise FuseOSError(ENOENT)
            if not isinstance(size, int):
                raise FuseOSError(ENOENT)
            size = max(size, 0)
            r = dict(st_mode=S_IFREG, st_nlink=1, st_size=size, st_blocks=size*2)
            r.update(self.common_attrs)
            return r
        else:
            try:
                val = self.flat[path]
            except KeyError:
                raise FuseOSError(ENOENT)
            if isinstance(val, int):
                r = dict(st_mode=S_IFDIR, st_nlink=2, st_size=0, st_blocks=0)
                r.update(self.common_attrs)
                return r
            else:
                r = dict(st_mode=S_IFDIR, st_nlink=2+len(val), st_size=0, st_blocks=0)
                r.update(self.common_attrs)
                return r
    

def main():
    parser = argparse.ArgumentParser(
        description='Show gnucash expenses as a filesystem, to analyze them with '
                    'disk usage analyzers.')
    parser.add_argument('fn', metavar='gnucash-filename',
                       help='filename of the gnucash file')
    parser.add_argument('mount', metavar='mountpoint',
                       help='directory to mount the filesystem on')

    args = parser.parse_args()
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    fuse = FUSE(Filesystem(args.fn), args.mount, foreground=True)

if __name__ == '__main__':
    main()

